package org.loong.crypto.extension.cms;

import java.io.IOException;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.CRLException;
import java.security.cert.CertificateEncodingException;
import java.security.cert.X509Certificate;
import java.util.Arrays;
import java.util.Collections;
import java.util.Enumeration;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.stream.Collectors;

import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.bouncycastle.asn1.ASN1Encoding;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.ASN1OctetString;
import org.bouncycastle.asn1.ASN1Sequence;
import org.bouncycastle.asn1.ASN1Set;
import org.bouncycastle.asn1.cms.CMSObjectIdentifiers;
import org.bouncycastle.asn1.cms.ContentInfo;
import org.bouncycastle.asn1.cms.EncryptedContentInfo;
import org.bouncycastle.asn1.cms.EnvelopedData;
import org.bouncycastle.asn1.cms.IssuerAndSerialNumber;
import org.bouncycastle.asn1.cms.SignerInfo;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.SubjectKeyIdentifier;
import org.bouncycastle.cert.jcajce.JcaCRLStore;
import org.bouncycastle.cert.jcajce.JcaCertStore;
import org.bouncycastle.cert.jcajce.JcaX509CertificateHolder;
import org.bouncycastle.cms.CMSEnvelopedData;
import org.bouncycastle.cms.CMSEnvelopedDataGenerator;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.cms.CMSProcessableByteArray;
import org.bouncycastle.cms.CMSSignedData;
import org.bouncycastle.cms.CMSSignedDataGenerator;
import org.bouncycastle.cms.CMSTypedData;
import org.bouncycastle.cms.DefaultSignedAttributeTableGenerator;
import org.bouncycastle.cms.SignerId;
import org.bouncycastle.cms.SignerInfoGenerator;
import org.bouncycastle.cms.SignerInformation;
import org.bouncycastle.cms.SignerInformationStore;
import org.bouncycastle.cms.jcajce.JcaSignerInfoGeneratorBuilder;
import org.bouncycastle.operator.BufferingContentSigner;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.OperatorCreationException;
import org.bouncycastle.operator.OperatorException;
import org.bouncycastle.operator.OutputEncryptor;
import org.loong.crypto.common.exception.CryptoException;
import org.loong.crypto.core.algorithm.CipherAlgorithm;
import org.loong.crypto.core.algorithm.EncryptionMethod;
import org.loong.crypto.core.algorithm.MessageDigestAlgorithm;
import org.loong.crypto.core.algorithm.SignatureAlgorithm;
import org.loong.crypto.core.params.CipherParameters;
import org.loong.crypto.core.params.SignatureParameters;
import org.loong.crypto.core.utils.ECKeyUtils;
import org.loong.crypto.core.utils.KeyTools;
import org.loong.crypto.extension.cms.jce.JceAsymmetricKeyUnwrapper;
import org.loong.crypto.extension.cms.jce.JceCMSContentEncryptorBuilder;
import org.loong.crypto.extension.cms.jce.JceKeyTransRecipientInfoGenerator;
import org.loong.crypto.extension.operator.jce.JceContentSignerBuilder;
import org.loong.crypto.extension.operator.jce.JceDigestCalculatorProviderBuilder;
import org.loong.crypto.model.asn1.cms.SignedAndEnvelopedData;
import org.loong.crypto.service.CryptoService;
import org.loong.crypto.spi.Join;

/**
 * The signedAndEnvelopedData provider.
 */
@Join
public class SignedAndEnvelopedDataProvider implements CMSProvider {

    private final EnvelopedDataHelper helper = new EnvelopedDataHelper();

    @Override
    public byte[] generate(CryptoService cryptoService, CMSParameters params) throws CryptoException {
        if (ArrayUtils.isEmpty(params.getData())) {
            throw new CryptoException("待加密的数据不允许为空.");
        }

        if (Objects.isNull(params.getEncryptionMethod())) {
            throw new CryptoException("数据加密算法不允许为空.");
        }

        final X509Certificate encryptionCertificate = params.getEncryptionCertificate();
        final PublicKey encryptionPublicKey = params.getEncryptionPublicKey();
        if (Objects.isNull(encryptionCertificate) && Objects.isNull(encryptionPublicKey)) {
            throw new CryptoException("接收者证书或公钥不允许为空.");
        }

        if (Objects.isNull(params.getSignatureAlgorithm())) {
            throw new CryptoException("数据签名算法不允许为空.");
        }

        final X509Certificate signingCertificate = params.getSigningCertificate();
        final PublicKey signingPublicKey = params.getSigningPublicKey();
        if (Objects.isNull(signingCertificate) && Objects.isNull(signingPublicKey)) {
            throw new CryptoException("签名者证书或公钥不允许为空.");
        }

        final PrivateKey privateKey = params.getPrivateKey();
        if (Objects.isNull(privateKey)) {
            throw new CryptoException("签名者私钥不允许为空.");
        }

        try {
            CMSEnvelopedDataGenerator envelopedDataGenerator = new CMSEnvelopedDataGenerator();
            JceKeyTransRecipientInfoGenerator recipientGenerator = null;
            if (Objects.nonNull(encryptionCertificate)) {
                recipientGenerator = new JceKeyTransRecipientInfoGenerator(cryptoService, encryptionCertificate);
            } else if (Objects.nonNull(encryptionPublicKey)) {
                byte[] subjectKeyId = KeyTools.createSubjectKeyId(encryptionPublicKey).getKeyIdentifier();
                recipientGenerator = new JceKeyTransRecipientInfoGenerator(cryptoService, subjectKeyId, encryptionPublicKey);
            }
            envelopedDataGenerator.addRecipientInfoGenerator(recipientGenerator);

            final SignatureAlgorithm signatureAlgorithm = SignatureAlgorithm.find(params.getSignatureAlgorithm());
            ASN1ObjectIdentifier contentType = CMSObjectIdentifiers.data;
            if (ECKeyUtils.isSM2Algorithm(signatureAlgorithm.getOID())) {
                contentType = GMCMSObjectIdentifiers.data;
            }

            final CMSTypedData cmsData = new CMSProcessableByteArray(contentType, params.getData());
            final OutputEncryptor encryptor = new JceCMSContentEncryptorBuilder(cryptoService, EncryptionMethod.find(params.getEncryptionMethod())).build();
            final CMSEnvelopedData cmsEnvelopedData = envelopedDataGenerator.generate(cmsData, encryptor);
            final EnvelopedData envelopedData = EnvelopedData.getInstance(cmsEnvelopedData.toASN1Structure().getContent());

            final JceContentSignerBuilder contentSignerBuilder = new JceContentSignerBuilder(cryptoService, signatureAlgorithm.getName());
            final ContentSigner signer = new BufferingContentSigner(contentSignerBuilder.build(privateKey), 20480);
            final JcaSignerInfoGeneratorBuilder signerInfoGeneratorBuilder = new JcaSignerInfoGeneratorBuilder(new JceDigestCalculatorProviderBuilder(cryptoService).build())
                    .setSignedAttributeGenerator(new DefaultSignedAttributeTableGenerator()).setDirectSignature(params.isHasNoSignedAttributes());
            SignerInfoGenerator signerInfoGenerator = null;
            SignerId signerId = null;
            if (Objects.nonNull(signingCertificate)) {
                signerInfoGenerator = signerInfoGeneratorBuilder.build(signer, signingCertificate);
                final X500Name issuer = X500Name.getInstance(signingCertificate.getIssuerX500Principal().getEncoded());
                signerId = new SignerId(issuer, signingCertificate.getSerialNumber());
            } else if (Objects.nonNull(signingPublicKey)) {
                final byte[] subjectKeyId = KeyTools.createSubjectKeyId(signingPublicKey).getKeyIdentifier();
                signerInfoGenerator = signerInfoGeneratorBuilder.build(signer, subjectKeyId);
                signerId = new SignerId(subjectKeyId);
            }

            final CMSSignedDataGenerator signedDataGenerator = new CMSSignedDataGenerator();
            signedDataGenerator.addSignerInfoGenerator(signerInfoGenerator);

            if (Objects.nonNull(signingCertificate)) {
                signedDataGenerator.addCertificates(new JcaCertStore(Collections.singletonList(signingCertificate)));
            }

            if (ArrayUtils.isNotEmpty(params.getCrls())) {
                signedDataGenerator.addCRLs(new JcaCRLStore(Arrays.asList(params.getCrls())));
            }

            final CMSSignedData cmsSignedData = signedDataGenerator.generate(cmsData, true);

            final SignerInformationStore signerStore = cmsSignedData.getSignerInfos();
            final SignerInformation signerInformation = signerStore.get(signerId);
            final AlgorithmIdentifier digestAlgorithm = signerInformation.getDigestAlgorithmID();
            final ASN1Set digestAlgorithms = CMSUtils.createBERSetFromList(Collections.singletonList(digestAlgorithm));

            final ASN1Set certificates = Optional.ofNullable(signingCertificate).map(certificate -> CMSUtils.fetchCertificatesList(certificate)).orElse(null);
            final ASN1Set crls = null;

            final List<SignerInfo> signers = signerStore.getSigners().stream().map(signerInfo -> signerInfo.toASN1Structure()).collect(Collectors.toList());
            final ASN1Set signerInfos = CMSUtils.createBERSetFromList(signers);

            contentType = CMSObjectIdentifiers.signedAndEnvelopedData;
            if (ECKeyUtils.isSM2Algorithm(signatureAlgorithm.getOID())) {
                contentType = GMCMSObjectIdentifiers.signedAndEnvelopedData;
            }

            return new ContentInfo(contentType,
                    new SignedAndEnvelopedData(envelopedData.getRecipientInfos(), digestAlgorithms, envelopedData.getEncryptedContentInfo(), certificates, crls, signerInfos))
                            .getEncoded(ASN1Encoding.DER);
        } catch (CertificateEncodingException | CMSException | NoSuchAlgorithmException | OperatorCreationException | IOException | CRLException e) {
            throw new CryptoException(e.getMessage(), e);
        }
    }

    @Override
    public byte[] parse(CryptoService cryptoService, CMSParameters params) throws CryptoException {
        if (ArrayUtils.isEmpty(params.getData())) {
            throw new CryptoException("待解密的数据不允许为空.");
        }

        final X509Certificate encryptionCertificate = params.getEncryptionCertificate();
        final PublicKey encryptionPublicKey = params.getEncryptionPublicKey();
        if (Objects.isNull(encryptionCertificate) && Objects.isNull(encryptionPublicKey)) {
            throw new CryptoException("接收者证书或公钥不允许为空.");
        }

        final X509Certificate signingCertificate = params.getSigningCertificate();
        final PublicKey signingPublicKey = params.getSigningPublicKey();
        if (Objects.isNull(signingCertificate) && Objects.isNull(signingPublicKey)) {
            throw new CryptoException("签名者证书或公钥不允许为空.");
        }

        final PrivateKey privateKey = params.getPrivateKey();
        if (Objects.isNull(privateKey)) {
            throw new CryptoException("接收者私钥不允许为空.");
        }

        try {
            final SignedAndEnvelopedData signedAndEnvelopedData = SignedAndEnvelopedData.getInstance(CMSUtils.fetchWrappedData(params.getData()));
            if (Objects.isNull(signedAndEnvelopedData)) {
                return null;
            }

            final Enumeration<?> recipientInfos = signedAndEnvelopedData.getRecipientInfos().getObjects();
            Pair<AlgorithmIdentifier, byte[]> encryptedKeyData = null;
            if (Objects.nonNull(encryptionCertificate)) {
                encryptedKeyData = CMSUtils.fetchEncryptedKeyDatas(encryptionCertificate, recipientInfos);
            } else if (Objects.nonNull(encryptionPublicKey)) {
                encryptedKeyData = CMSUtils.fetchEncryptedKeyDatas(encryptionPublicKey, recipientInfos);
            }

            if (Objects.isNull(encryptedKeyData)) {
                throw new CryptoException("密钥加密算法或数据加密密钥为空.");
            }

            AlgorithmIdentifier keyEncryptionAlgorithm = encryptedKeyData.getLeft();
            byte[] encryptedKey = encryptedKeyData.getRight();

            final JceAsymmetricKeyUnwrapper unwrapper = new JceAsymmetricKeyUnwrapper(keyEncryptionAlgorithm, cryptoService, privateKey);
            final EncryptedContentInfo encryptedContentInfo = signedAndEnvelopedData.getEncryptedContentInfo();
            final AlgorithmIdentifier contentEncryptionAlgorithm = encryptedContentInfo.getContentEncryptionAlgorithm();
            final Key key = helper.getKey(contentEncryptionAlgorithm.getAlgorithm(), unwrapper.generateUnwrappedKey(contentEncryptionAlgorithm, encryptedKey));
            final byte[] iv = helper.getIv(key, contentEncryptionAlgorithm);

            final byte[] content = cryptoService.decrypt(CipherAlgorithm.find(helper.getBaseCipherName(contentEncryptionAlgorithm.getAlgorithm())),
                    CipherParameters.builder().encryptionMethod(helper.getEncryptionMethod(contentEncryptionAlgorithm.getAlgorithm())).key(key).iv(iv).build(),
                    encryptedContentInfo.getEncryptedContent().getOctets());

            IssuerAndSerialNumber issuerAndSerial = null;
            byte[] signingKeyId = null;
            SignatureParameters signatureParameters = null;
            if (Objects.nonNull(signingCertificate)) {
                issuerAndSerial = new IssuerAndSerialNumber(new JcaX509CertificateHolder(signingCertificate).toASN1Structure());
                signingKeyId = KeyTools.createSubjectKeyId(signingCertificate.getPublicKey()).getKeyIdentifier();
                signatureParameters = SignatureParameters.builder().key(signingCertificate.getPublicKey()).build();
            } else if (Objects.nonNull(signingPublicKey)) {
                signingKeyId = KeyTools.createSubjectKeyId(signingPublicKey).getKeyIdentifier();
                signatureParameters = SignatureParameters.builder().key(signingPublicKey).build();
            }

            boolean signValid = false;
            final ASN1Set signerInfos = signedAndEnvelopedData.getSignerInfos();
            for (int i = 0; i < signerInfos.size() && Boolean.FALSE.equals(signValid); i++) {
                final SignerInfo signerInfo = SignerInfo.getInstance(signerInfos.getObjectAt(i));

                if (Objects.nonNull(issuerAndSerial)) {
                    IssuerAndSerialNumber ias = IssuerAndSerialNumber.getInstance(signerInfo.getSID());
                    if (!Objects.equals(ias, issuerAndSerial)) {
                        continue;
                    }
                } else if (Objects.nonNull(signingKeyId)) {
                    SubjectKeyIdentifier ski = SubjectKeyIdentifier.getInstance(signerInfo.getSID().getId());
                    if (!Arrays.equals(signingKeyId, ski.getKeyIdentifier())) {
                        continue;
                    }
                }

                final SignatureAlgorithm signatureAlgorithm = SignatureAlgorithm.findByAlgorithmId(signerInfo.getDigestEncryptionAlgorithm());
                final byte[] encryptedDigest = signerInfo.getEncryptedDigest().getOctets();

                ASN1Set authenticatedAttributes = signerInfo.getAuthenticatedAttributes();
                if (Objects.nonNull(authenticatedAttributes)) {
                    MessageDigestAlgorithm digestAlgorithm = MessageDigestAlgorithm.findByAlgorithmId(signerInfo.getDigestAlgorithm());
                    byte[] digest = cryptoService.digest(digestAlgorithm, content);

                    for (int j = 0; j < authenticatedAttributes.size(); j++) {
                        ASN1Sequence authAttrSeq = ASN1Sequence.getInstance(authenticatedAttributes.getObjectAt(j));
                        ASN1ObjectIdentifier attrOid = ASN1ObjectIdentifier.getInstance(authAttrSeq.getObjectAt(0));
                        if (PKCSObjectIdentifiers.pkcs_9_at_messageDigest.equals(attrOid)) {
                            ASN1Set messageDigest = ASN1Set.getInstance(authAttrSeq.getObjectAt(1));
                            ASN1OctetString asn1ObjString = ASN1OctetString.getInstance(messageDigest.getObjectAt(0));
                            if (Objects.deepEquals(digest, asn1ObjString.getOctets())) {
                                signValid = cryptoService.verify(signatureAlgorithm, signatureParameters, authenticatedAttributes.getEncoded(ASN1Encoding.DER), encryptedDigest);
                            }
                        }
                    }
                } else {
                    signValid = cryptoService.verify(signatureAlgorithm, signatureParameters, content, encryptedDigest);
                }
            }

            if (Boolean.FALSE.equals(signValid)) {
                throw new CryptoException("Signature validation failed.");
            }
            return content;
        } catch (IOException | CertificateEncodingException | CMSException | OperatorException | NoSuchAlgorithmException e) {
            throw new CryptoException(e.getMessage(), e);
        }
    }
}
